import cv2 as cv
import numpy as np
from utils.transition import linearFloat, linear

class PitchV():
  def __init__(self, outputPath, size=(610, 1080), fps=24):
    fourcc = cv.VideoWriter_fourcc(*'MP4V')
    self.__wait = int(1000 / 24)
    self.__vw = cv.VideoWriter(outputPath, fourcc, fps, size)

  # 淡入淡出的形式切换图片
  def opacityPitch(self, bg, front, totalFrames, transformFrames):
    bh, bw = bg.shape[:2]
    nh, nw = front.shape[:2]
    if bh != nh and bw != nw:
      raise ValueError('图片尺寸不符')

    if transformFrames == 0:
      self.__waitFames(front, totalFrames)  
    else:  
      steps = linearFloat(transformFrames, 0, 1)
      
      for i, alpha in enumerate(steps):
        beta = 1 - alpha
        gamma = 0
        res = cv.addWeighted(front, alpha, bg, beta, gamma)
        self.__vw.write(res)
        # cv.imshow('frame', res)
        cv.waitKey(self.__wait)

      restFrames = totalFrames - transformFrames
      self.__waitFames(front, restFrames)  


  # 从右往左的形式切换图片
  def unfoldPitch(self, bg, front, totalFrames, transformFrames):
    bh, bw = bg.shape[:2]
    nh, nw = front.shape[:2]
    if bh != nh and bw != nw:
      raise ValueError('图片尺寸不符')

    screenW = 5
    steps = linear(transformFrames, 0, bw)
    screen = np.ones((bh, screenW, 3), dtype=np.uint8) * 240

    for i, item in enumerate(steps):
      if i == 0:
        # cv.imshow('frame', bg)
        self.__vw.write(bg)
        cv.waitKey(self.__wait)
        continue

      tlx = int(bw-item)
      trx = bw
      fig = front[: , 0:item]
      if bw-item >= screenW:
        bg[:, int(bw-item-screenW):int(bw-item)] = screen
      else:
        wscreen = np.ones((bh, bw, 3), dtype=np.uint8) * 240
        bg = wscreen
      fig = cv.resize(fig, ((trx-tlx), bh))
      bg[:, tlx:trx] = fig
      # cv.imshow('frame', bg)
      self.__vw.write(bg)
      cv.waitKey(self.__wait)

    restFrames = totalFrames - transformFrames
    self.__waitFames(front, waitFrame=restFrames) 

  # 图片停留时间
  def __waitFames(self, img, waitFrame):
    for i in range(0, waitFrame):
      # cv.imshow('frame', img)
      self.__vw.write(img)
      cv.waitKey(self.__wait)
